import base64
import os
import re
import sys

CURRENT_PATH = os.path.dirname(os.path.abspath(__file__))


def transPicToBase64():
    img_dir = os.path.join(CURRENT_PATH, 'img')
    pic_list = os.listdir(img_dir)
    for pic in pic_list:
        pic_path = os.path.join("img", pic)
        ext = pic.split(".")[-1]
        with open(pic_path, 'rb') as pic_file:
            img_content = pic_file.read()
        output = base64.b64encode(img_content).decode()
        pic_name = ".".join(pic.split(".")[0:-1])
        base_file = pic_name + 'Base64.txt'

        with open(base_file, 'a') as file:
            file.write("data:image/{ext};base64,{data}".format(ext=ext, data=output))


def transBase64ToPic():
    base64_dir = os.path.join(CURRENT_PATH, 'base64')
    base64_list = os.listdir(base64_dir)
    for text in base64_list:
        # 二进制方式打开图文件
        file_path = os.path.join("base64", text)
        with open(file_path, 'rb') as base64_file:
            base64_content = base64_file.read().decode()
            print(type(base64_content))
        result = re.search("data:image/(?P<ext>.*?);base64,(?P<data>.*)", base64_content, re.DOTALL)
        if result:
            ext = result.groupdict().get("ext")
            data = result.groupdict().get("data")
        else:
            print("The data format is invalid")
            return

        img_content = base64.b64decode(data)
        filename = "{}.{}".format(text.split(".")[0], ext)
        with open(filename, "wb") as f:
            f.write(img_content)


if __name__ == "__main__":
    action = sys.argv[1]
    if action == "b2p":
        transBase64ToPic()
    elif action == "p2b":
        transPicToBase64()
    else:
        LOGGER.error("Invalid action %s. Available options are as follows:[p2b,b2p]",
                     sys.argv[1])
